from cryptography.fernet import Fernet
from typing import Union, Optional
from pathlib import Path


class EncryptionManager:
    """
    AES-256加密管理器，提供数据和文件的加密解密功能

    使用示例:
    >>> manager = EncryptionManager()
    >>> encrypted = manager.encrypt_data("secret")
    >>> decrypted = manager.decrypt_data(encrypted)
    """

    def __init__(self, key_dir: Union[str, Path] = "keys"):
        """
        初始化加密管理器
        :param key_dir: 密钥存储目录
        """
        self.key_dir = Path(key_dir)
        self.key_dir.mkdir(exist_ok=True, parents=True)
        self.key_path = self.key_dir / "encryption.key"
        self.key = self._load_or_generate_key()

    def _load_or_generate_key(self) -> bytes:
        """加载或生成加密密钥"""
        if self.key_path.exists():
            with open(self.key_path, "rb") as key_file:
                return key_file.read()

        key = Fernet.generate_key()
        with open(self.key_path, "wb") as key_file:
            key_file.write(key)
        return key

    def encrypt_data(self, data: Union[str, bytes]) -> bytes:
        """
        加密数据
        :param data: 要加密的数据（字符串或字节）
        :return: 加密后的字节数据
        """
        if isinstance(data, str):
            data = data.encode('utf-8')
        return Fernet(self.key).encrypt(data)

    def decrypt_data(self, encrypted_data: bytes) -> str:
        """
        解密数据
        :param encrypted_data: 加密的数据
        :return: 解密后的字符串
        """
        return Fernet(self.key).decrypt(encrypted_data).decode('utf-8')

    def encrypt_file(
            self,
            input_path: Union[str, Path],
            output_path: Optional[Union[str, Path]] = None,
            chunk_size: int = 64 * 1024
    ) -> Path:
        """
        加密文件
        :param input_path: 输入文件路径
        :param output_path: 输出文件路径（可选）
        :param chunk_size: 分块大小（字节）
        :return: 加密后的文件路径
        """
        input_path = Path(input_path)
        if not input_path.exists():
            raise FileNotFoundError(f"输入文件不存在: {input_path}")

        if output_path is None:
            output_path = input_path.parent / f"{input_path.name}.enc"
        output_path = Path(output_path)

        cipher = Fernet(self.key)
        with open(input_path, "rb") as infile, open(output_path, "wb") as outfile:
            while chunk := infile.read(chunk_size):
                encrypted = cipher.encrypt(chunk)
                outfile.write(encrypted)

        return output_path

    def decrypt_file(
            self,
            input_path: Union[str, Path],
            output_path: Optional[Union[str, Path]] = None,
            chunk_size: int = 64 * 1024
    ) -> Path:
        """
        解密文件
        :param input_path: 加密文件路径
        :param output_path: 输出文件路径（可选）
        :param chunk_size: 分块大小（字节）
        :return: 解密后的文件路径
        """
        input_path = Path(input_path)
        if not input_path.exists():
            raise FileNotFoundError(f"输入文件不存在: {input_path}")

        if output_path is None:
            if input_path.suffix == ".enc":
                output_path = input_path.with_suffix("")
            else:
                output_path = input_path.parent / f"{input_path.name}.decrypted"
        output_path = Path(output_path)

        cipher = Fernet(self.key)
        with open(input_path, "rb") as infile, open(output_path, "wb") as outfile:
            while chunk := infile.read(chunk_size):
                decrypted = cipher.decrypt(chunk)
                outfile.write(decrypted)

        return output_path